﻿using System;
using System.Collections;
using System.Diagnostics;
using System.Linq;
using System.Windows;
using System.Windows.Controls;
using System.Collections.Generic;
using System.Windows.Data;

namespace VirtualizingTreeView.Control
{
	public enum SelectionMode
	{
		Single,
		Multiple
	}

	public class VirtualizingTree: ItemsControl
	{
		private readonly BoolToVisibilityConverter _boolToVisibilityConverter;
		private Boolean _suppressCallback;
		private VirtualizingDataContainer _selectedDataContainer;

		public event EventHandler<VirtualizingSelectionEventArgs> SelectionChanged;
		public event EventHandler<VirtualizingEventArgs> Expanded;

		public static readonly DependencyProperty IndentProperty = DependencyProperty.Register(
			"Indent",
			typeof( Double ),
			typeof( VirtualizingTree ),
			null );

		public static readonly DependencyProperty SelectionModeProperty = DependencyProperty.Register(
			"SelectionMode",
			typeof (SelectionMode),
			typeof (VirtualizingTree),
			new PropertyMetadata( SelectionMode.Single ) );

		public static readonly DependencyProperty SelectedItemProperty = DependencyProperty.Register(
			"SelectedItem",
			typeof (Object),
			typeof (VirtualizingTree),
			new PropertyMetadata( null, OnSelectedItemChanged ) );

		public static readonly DependencyProperty MultiSelectItemStyleProperty = DependencyProperty.Register(
			"MultiSelectItemStyle",
			typeof (Style),
			typeof (VirtualizingTree),
			null );

		public static readonly DependencyProperty ExpandedItemTemplateProperty = DependencyProperty.Register(
			"ExpandedItemTemplate",
			typeof (DataTemplate),
			typeof (VirtualizingTree),
			null );

		public static readonly DependencyProperty HasChildrenPathProperty = DependencyProperty.Register(
			"HasChildrenPath",
			typeof( String ),
			typeof( VirtualizingTree ),
			null );

		public static readonly DependencyProperty IsSelectedPathProperty = DependencyProperty.Register(
			"IsSelectedPath",
			typeof (String),
			typeof (VirtualizingTree),
			null );

		public Object SelectedItem
		{
			get { return GetValue( SelectedItemProperty ); }
			set { SetValue( SelectedItemProperty, value );}
		}

		public Double Indent
		{
			get { return ( Double )GetValue( IndentProperty ); }
			set { SetValue( IndentProperty, value ); }
		}

		public SelectionMode SelectionMode
		{
			get { return ( SelectionMode )GetValue( SelectionModeProperty ); }
			set { SetValue( SelectionModeProperty, value );}
		}

		public Style MultiSelectItemStyle
		{
			get { return ( Style ) GetValue( MultiSelectItemStyleProperty ); }
			set { SetValue( MultiSelectItemStyleProperty, value );}
		}

		public DataTemplate ExpandedItemTemplate
		{
			get { return ( DataTemplate ) GetValue( ExpandedItemTemplateProperty ); }
			set { SetValue( ExpandedItemTemplateProperty, value );}
		}

		public String HasChildrenPath
		{
			get { return ( String )GetValue( HasChildrenPathProperty ); }
			set { SetValue( HasChildrenPathProperty, value );}
		}

		private VirtualizingDataContainer SelectedDataContainer
		{
			get { return _selectedDataContainer; }
			set
			{
				_selectedDataContainer = value;
				SetValueNoCallback( SelectedItemProperty, _selectedDataContainer == null ? null : _selectedDataContainer.BoundItem );
			}
		}

		public VirtualizingTree()
		{
			DefaultStyleKey = typeof( VirtualizingTree );

			_boolToVisibilityConverter = new BoolToVisibilityConverter();
		}

		protected override DependencyObject GetContainerForItemOverride()
		{
			VirtualizingTreeItem treeItem = new VirtualizingTreeItem();
			if(HasChildrenPath != null)
			{
				treeItem.SetBinding( VirtualizingTreeItem.ExpanderVisibilityProperty,
					new Binding
					{
						Converter = _boolToVisibilityConverter,
						Path = new PropertyPath(HasChildrenPath)
					} );
			}

			
			return treeItem;
		}

		protected override Boolean IsItemItsOwnContainerOverride( Object item )
		{
			return (item is VirtualizingTreeItem);
		}

		protected override void PrepareContainerForItemOverride( DependencyObject element, Object item )
		{
			VirtualizingDataContainer dataContainer = item as VirtualizingDataContainer;
			if(dataContainer != null)
			{
				VirtualizingTreeItem treeItem = element as VirtualizingTreeItem;
				if( treeItem != null )
				{
					if(SelectionMode == SelectionMode.Multiple && MultiSelectItemStyle != null)
					{
						treeItem.Style = MultiSelectItemStyle;
					}

					Int32 nestingLevel = dataContainer.NestingLevel;
					treeItem.Margin = new Thickness( 
						nestingLevel * Indent, 
						treeItem.Margin.Top, 
						treeItem.Margin.Right, 
						treeItem.Margin.Bottom );

					if(ItemTemplate != null)
					{
					    treeItem.ContentTemplate = ItemTemplate;
					}

					if(ExpandedItemTemplate != null)
					{
						treeItem.ExpandedContentTemplate = ExpandedItemTemplate;
					}

					treeItem.Content = dataContainer;
				}
			}	
		}

		public void Add( IEnumerable items )
		{
			AddInternal( null, items );
		}

		public void Add( Object parentItem, IEnumerable items)
		{
			if(parentItem != null)
			{
				VirtualizingDataContainer dataContainer = GetDataContainer( parentItem );
				AddInternal( dataContainer, items );	
			}
			else
			{
				AddInternal( null, items );
			}
		}

		private void AddInternal( VirtualizingDataContainer dataContainer, IEnumerable items )
		{
			if(items == null)
			{
				throw new ArgumentNullException("items");
			}

			if( dataContainer != null )
			{
				Int32 index = Items.IndexOf( dataContainer ) + 1;
				if(index == 0)
				{
					throw new ArgumentException( "!Items.Contains(dataContainer)", "dataContainer" );
				}

				Int32 nestingLevel = dataContainer.NestingLevel + 1;

				foreach ( Object item in items )
				{
					Items.Insert( index, new VirtualizingDataContainer( item, nestingLevel ) );
					index++;
				}
			}
			else
			{
				foreach( Object item in items )
				{
					Items.Add( new VirtualizingDataContainer( item, 0 ) );
				}
			}
		}

		public void Remove(Object item)
		{
			if(item == null)
			{
				throw new ArgumentNullException("item");
			}

			VirtualizingDataContainer container = GetDataContainer( item );
			if(container != null)
			{
				RemoveInternal( container );	
			}
		}

		public void Collapse( Object item )
		{
			if(item == null)
			{
				throw new ArgumentNullException( "item" );
			}

			VirtualizingDataContainer container = GetDataContainer( item );
			if(container.IsExpanded)
			{
				VirtualizingTreeItem treeItem = GetTreeItem( container );
				if(treeItem != null)
				{
					treeItem.ExpandOrCollapse();
				}
				else
				{
					container.IsExpanded = false;
					OnCollapsedInternal(container);
				}
			}
		}

		public Boolean Contains( Object item )
		{
			if(item == null)
			{
				return false;
			}

			return GetDataContainer( item ) != null;
		}

		public Boolean HasChildren( Object item )
		{
			VirtualizingDataContainer dataContainer = GetDataContainer( item );
			Int32 index = Items.IndexOf( dataContainer );
			if(index < 0)
			{
				return false;
			}

			index++;
			if( index == Items.Count)
			{
				return false;
			}

			VirtualizingDataContainer nextDataContainer = ( VirtualizingDataContainer ) Items[ index ];
			return nextDataContainer.NestingLevel == dataContainer.NestingLevel + 1;
		}

		public Boolean IsExpanded( Object item )
		{
			if(item == null)
			{
				return false;
			}

			VirtualizingDataContainer dataContainer = GetDataContainer( item );
			return dataContainer.IsExpanded;
		}

		private void RemoveInternal( VirtualizingDataContainer dataContaier )
		{
			OnCollapsedInternal( dataContaier );
			Items.Remove( dataContaier );
		}

		private VirtualizingDataContainer GetParentDataContainer(VirtualizingDataContainer dataContainer)
		{
			if( dataContainer.NestingLevel == 0 )
			{
				return null;
			}

			Int32 index = Items.IndexOf( dataContainer ) - 1;
			while( index >= 0 )
			{
				VirtualizingDataContainer prevDataContainer = ( VirtualizingDataContainer )Items[index];
				if( prevDataContainer.NestingLevel < dataContainer.NestingLevel )
				{
					return prevDataContainer;
				}

				index--;
			}

			Debug.Assert( false, "Invalid tree structure" );
			return null;
		}

		private IEnumerable<VirtualizingDataContainer> GetChildren(VirtualizingDataContainer dataContainer)
		{
			Int32 index = Items.IndexOf( dataContainer ) + 1;
			while( index < Items.Count )
			{
				VirtualizingDataContainer nextDataContainer = ( VirtualizingDataContainer )Items[index];
				if( nextDataContainer.NestingLevel > dataContainer.NestingLevel )
				{
					yield return nextDataContainer;
				}
				else
				{
					break;					
				}
				index++;
			}
		}

		public void UpdateSelectionState(Object item)
		{
			VirtualizingDataContainer dc = GetDataContainer( item );
			UpdateChildrenSelectionState( dc );
			UpdateParentSelectionState( dc );
			
		}

		private void UpdateParentSelectionState( VirtualizingDataContainer item )
		{
			if(item == null)
			{
				return;
			}

			VirtualizingDataContainer parent = GetParentDataContainer( item );
			if(parent != null)
			{
				IEnumerator<VirtualizingDataContainer> children = GetChildren( parent ).GetEnumerator();
				children.MoveNext();

				Boolean? childrenCheckedState = children.Current.IsSelected;
				while ( children.MoveNext() )
				{
					if(childrenCheckedState != children.Current.IsSelected)
					{
						childrenCheckedState = null;
						break;
					}
				}

				SetSelection( parent, childrenCheckedState );
				UpdateParentSelectionState( parent );
			}
		}

		private void UpdateChildrenSelectionState( VirtualizingDataContainer item )
		{
			if( item.IsSelected.HasValue)
			{
				IEnumerator<VirtualizingDataContainer> children = GetChildren( item ).GetEnumerator();
				while ( children.MoveNext() )
				{
					SetSelection( children.Current, item.IsSelected );	
					UpdateChildrenSelectionState( children.Current );
				}
			}
		}

		public void Select(IEnumerable items)
		{
			if(items == null)
			{
				throw new ArgumentNullException("items");
			}

			SelectInternal( from Object item in items
							select GetDataContainer( item ) );
		}

		public void Unselect( IEnumerable items )
		{
			if( items == null )
			{
				throw new ArgumentNullException( "items" );
			}

			UnselectInternal( from Object item in items
							  select GetDataContainer( item ) );
		}


		private void SelectInternal(IEnumerable<VirtualizingDataContainer> dataContainers)
		{
			foreach( VirtualizingDataContainer dataContainer in dataContainers )
			{	
				if(dataContainer != null)
				{
					SetSelection( dataContainer, true );	
				}
				
			}
		}

		private void UnselectInternal( IEnumerable<VirtualizingDataContainer> dataContainers )
		{
			foreach( VirtualizingDataContainer dataContainer in dataContainers )
			{
				if(dataContainer != null)
				{
					SetSelection( dataContainer, false );	
				}
			}
		}

		private void SetSelection( VirtualizingDataContainer dataContainer, Boolean? b)
		{
			VirtualizingTreeItem treeItem = GetTreeItem( dataContainer );
			if( treeItem != null )
			{
				treeItem.Select(b);
			}
			else
			{
				dataContainer.IsSelected = b;
			}
		}

		internal void OnSelectInternal( VirtualizingDataContainer dataContainer )
		{
			Debug.Assert( dataContainer.IsSelected == true );

			if( SelectionMode == SelectionMode.Single )
			{
				if( SelectedDataContainer != null )
				{
					VirtualizingTreeItem treeItem = GetTreeItem( SelectedDataContainer );
					if( treeItem != null )
					{
						treeItem.Select( false );
					}
					else
					{
						SelectedDataContainer.IsSelected = false;
					}
				}
				RaiseSelectionChanged( SelectedDataContainer, dataContainer );
			}
			else
			{
				RaiseSelectionChanged( null, dataContainer );
			}
			
			SelectedDataContainer = dataContainer;
		}


		internal void OnUnselectInternal( VirtualizingDataContainer dataContainer )
		{
			Debug.Assert( dataContainer.IsSelected == false );
			Debug.Assert( SelectionMode != SelectionMode.Single || SelectedDataContainer == dataContainer );

			if(SelectedDataContainer == dataContainer)
			{
				SelectedDataContainer = null;
			}

			if(SelectionMode == SelectionMode.Multiple)
			{
				RaiseSelectionChanged( dataContainer, null );
			}
		}

		internal void OnUnknownStateInternal( VirtualizingDataContainer dataContainer)
		{
			Debug.Assert( dataContainer.IsSelected == null );
			Debug.Assert( SelectionMode == SelectionMode.Multiple );

			RaiseSelectionChanged( dataContainer, dataContainer );
		}

		internal void OnExpandedInternal( VirtualizingDataContainer dataContainer)
		{
			Debug.Assert( dataContainer.IsExpanded );

			if(Expanded != null)
			{
				Expanded( this, new VirtualizingEventArgs( dataContainer.BoundItem ) );
			}
		}

		internal void OnCollapsedInternal( VirtualizingDataContainer dataContainer)
		{
			Debug.Assert( !dataContainer.IsExpanded );

			Int32 index = Items.IndexOf( dataContainer ) + 1;
			while ( index < Items.Count )
			{
				VirtualizingDataContainer nextDataContainer = ( VirtualizingDataContainer ) Items[ index ];
				if(nextDataContainer.NestingLevel > dataContainer.NestingLevel)
				{
					Items.RemoveAt( index );
				}
				else
				{
					break;
				}
			}
		}

		private VirtualizingTreeItem GetTreeItem( VirtualizingDataContainer dataContainer)
		{
			return ( VirtualizingTreeItem ) ItemContainerGenerator.ContainerFromItem( dataContainer );
		}

		private VirtualizingDataContainer GetDataContainer( Object item )
		{
			VirtualizingDataContainer result = (VirtualizingDataContainer)Items.Where( dataContainer => dataContainer.Equals( item ) ).FirstOrDefault();
			//if( result == null )
			//{
			//    throw new ArgumentNullException( "parentItem is not associated with data container", "item" );
			//}
			return result;
		}

		private void SetValueNoCallback( DependencyProperty dp, Object value)
		{
			try
			{
				_suppressCallback = true;
				SetValue( dp, value );
			}
			finally
			{
				_suppressCallback = false;
			}
		}

		private void RaiseSelectionChanged( VirtualizingDataContainer removed, VirtualizingDataContainer added )
		{
			if(SelectionChanged != null)
			{
				SelectionChanged(this, new VirtualizingSelectionEventArgs( 
					removed != null ? removed.BoundItem : null,  
					added != null ? added.BoundItem : null ));
			}
		}

		private static void OnSelectedItemChanged( DependencyObject d, DependencyPropertyChangedEventArgs e )
		{
			VirtualizingTree tree = ( VirtualizingTree ) d;
			if ( tree._suppressCallback )
			{
				return;
			}
			if ( e.NewValue != null )
			{
				tree.Select( new[] {e.NewValue} );
			}
			else if ( e.OldValue != null )
			{
				tree.Unselect( new[] {e.OldValue} );
			}
		}
	}
}
